package zhongbo.javaweb.instagram.util.dns;

import javax.naming.Context;
import javax.naming.NamingException;
import javax.naming.directory.Attribute;
import javax.naming.directory.Attributes;
import javax.naming.directory.DirContext;
import javax.naming.directory.InitialDirContext;
import java.util.*;
import java.util.regex.Pattern;

/**
 * 说明：dns解析工具类
 * <p>
 * #link https://blog.51cto.com/7090376/1728735
 *
 * @author 钟波notebook
 * @since 2020年08月16日 星期日 13:04:11
 */
public class DnsUtil {

    private static final String IPV4_BASIC_PATTERN_STRING =
            // initial first field, 1-255
            "(([1-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\\.){1}" +
                    // following 2 fields, 0-255 followed by .
                    "(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\\.){2}" +
                    // final field, 0-255
                    "([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])";

    private static final Pattern IPV4_PATTERN =
            Pattern.compile("^" + IPV4_BASIC_PATTERN_STRING + "$");

    /**
     * 获取域名所有IP
     *
     * @param domain     域名
     * @param dnsServers DNS服务器列表
     * @param timeout    请求超时
     * @param retryCount 重试次数
     * @return 解析ip结果列表
     */
    public static Set<String> getAllIp(String domain, String[] dnsServers, int timeout, int retryCount) {
        Set<String> ips = new HashSet<>();

        for (String dnsServer : dnsServers) {
            List<String> ipList;
            try {
                ipList = getDnsRecs(domain, dnsServer, new String[]{"A"}, timeout, retryCount);
            } catch (NamingException e) {
                continue;
            }
            ips.addAll(ipList);
        }

        return ips;
    }

    /**
     * 获取DNS服务器信息
     *
     * @param domain     要获取DNS信息的域名
     * @param provider   DNS服务器
     * @param types      信息类型 "A"(IP信息)，"MX"
     * @param timeout    请求超时
     * @param retryCount 重试次数
     * @return 所有信息组成的数组
     * @throws NamingException 域名解析失败异常
     */

    @SuppressWarnings("rawtypes")
    public static ArrayList<String> getDnsRecs(String domain, String provider, String[] types, int timeout, int retryCount) throws NamingException {

        ArrayList<String> results = new ArrayList<>(15);

        Hashtable<String, String> env = new Hashtable<>();

        env.put("java.naming.factory.initial", "com.sun.jndi.dns.DnsContextFactory");

        //设置域名服务器
        env.put(Context.PROVIDER_URL, "dns://" + provider);

        // 连接时间
        env.put("com.sun.jndi.dns.timeout.initial", String.valueOf(timeout));

        // 连接次数
        env.put("com.sun.jndi.dns.timeout.retries", String.valueOf(retryCount));

        DirContext initialDirContext = new InitialDirContext(env);
        Attributes attrs = initialDirContext.getAttributes(domain, types);

        for (Enumeration e = attrs.getAll(); e.hasMoreElements(); ) {
            Attribute a = (Attribute) e.nextElement();
            int size = a.size();
            for (int i = 0; i < size; i++) {
                results.add((String) a.get(i));
            }
        }
        return results;
    }

    /**
     * 判断是否ipv4地址
     * <p>
     * 引用自 org.apache.http.conn.util.InetAddressUtils
     *
     * @param host 地址str
     * @return 是否ipv4地址
     */
    public static boolean isIPv4Address(String host) {
        return IPV4_PATTERN.matcher(host).matches();
    }

}
